import numpy as np
import pandas as pd

def read_transformation_matrices(file_path):
    df = pd.read_csv(file_path)
    
    matrices = []
    for index, row in df.iterrows():
        # 提取旋转矩阵部分
        rotation = np.array([
            [row['Vx0'], row['Vy0'], row['Vz0']],
            [row['Vx1'], row['Vy1'], row['Vz1']],
            [row['Vx2'], row['Vy2'], row['Vz2']]
        ])
        
        # 提取平移向量部分
        translation = np.array([row['P0'], row['P1'], row['P2']]).reshape(3, 1)
        
        matrices.append({
            'rot': rotation,
            'tran': translation
        })
    
    return matrices

def save_optimal_transformation(R_opt, t_opt, file_path):
    # 将最优旋转矩阵和平移向量保存为 .npz 文件
    np.savez(file_path, R_opt=R_opt, t_opt=t_opt)

def load_optimal_transformation(file_path):
    # 从 .npz 文件加载最优旋转矩阵和平移向量
    data = np.load(file_path)
    R_opt = data['R_opt']
    t_opt = data['t_opt']
    return R_opt, t_opt



def main():
    # 替换为你的CSV文件路径
    file_path_A = 'foundation_pose_matrices1.csv'
    file_path_B = 'data/sim_matrice.csv'
    
    # 读取刚体变换矩阵
    A = read_transformation_matrices(file_path_A)
    B = read_transformation_matrices(file_path_B)

    n = len(A)

    mean_tA = np.mean([a['tran'] for a in A], axis=0)
    mean_tB = np.mean([b['tran'] for b in B], axis=0)

    print("mean_tA:", mean_tA)
    print("mean_tB:", mean_tB)

    
    tAi_centered = [a['tran'] - mean_tA for a in A]
    tBi_centered = [b['tran'] - mean_tB for b in B]
    
    H = np.zeros((3, 3))
    for tAi_c, tBi_c in zip(tAi_centered, tBi_centered):
        tBi_c = tBi_c.reshape(1, 3)
        H -= tAi_c @ tBi_c # 注意负号

    print("H:", H)

    M_rot = np.zeros((3, 3))
    for a, b in zip(A, B):
        Ai_rot = a['rot']
        Bi_rot = b['rot']
        M_rot += Bi_rot.T @ Ai_rot

    print("M_rot:", M_rot)

    # M = M_rot+H
    M = M_rot
    print("M:", M)

    U, S, Vt = np.linalg.svd(M)
    print("U:", U)
    print("S:", S)
    print("Vt:", Vt)

    R_opt =  Vt.T @ U.T
    print("R_opt:", R_opt)
    print("det(R_opt):", np.linalg.det(R_opt))
    
    t_opt = mean_tB - R_opt @ mean_tA

    print("t_opt:", t_opt)

    transformed_rotation = R_opt @ A[0]['rot']
    transformed_translation = R_opt @ A[0]['tran'] + t_opt
    print("transformed_rotation:", transformed_rotation)
    print("transformed_translation:", transformed_translation)

    print("B0 R:", B[0]['rot'])
    print("B0 t:", B[0]['tran'])

    save_optimal_transformation(R_opt, t_opt, "best_transformation.npz")
    print("最优变换矩阵已保存为 best_transformation.npz")

    

if __name__ == "__main__":
    main()